import math
from dataclasses import dataclass, field
import numpy as np


@dataclass
class Trajectory:
    tra_id: int = 0
    position_list: list = field(default_factory=list)
    # 记录运动的距离，考虑无人机自身的运动
    moving_distance_list: list = field(default_factory=list)
    # 记录目标对角线的长度
    diag_length_list: list = field(default_factory=list)
    diag_length_mean: float = 0
    diag_length_std: float = 0
    # 记录速度
    speed_list: float = field(default_factory=list)
    image_size: tuple = (1920, 1080)
    # 记录最后一帧的序号
    last_frame_index: int = 0

    def add_position(self, trk_box):
        trk_id, x1, y1, x2, y2 = trk_box
        assert int(trk_id.item()) == self.tra_id

        img_w, img_h = self.image_size
        x1 = int(max(min(x1.item(), 1), 0) * img_w)
        y1 = int(max(min(y1.item(), 1), 0) * img_h)
        x2 = int(max(min(x2.item(), 1), 0) * img_w)
        y2 = int(max(min(y2.item(), 1), 0) * img_h)
        self.position_list.append([x1, y1, x2, y2])

        alpha, lambda_ = 0.5, 2.55
        diag_length = math.sqrt((x2 - x1)**2 + (y2 - y1)**2)
        if self.diag_length_std == 0:  # 初始化
            self.diag_length_mean = diag_length
            self.diag_length_std = 20
        if diag_length < lambda_ * self.diag_length_std:
            self.diag_length_std = math.sqrt(
                alpha * self.diag_length_std**2 + \
                (1 - alpha) * (diag_length - self.diag_length_mean)**2)
            self.diag_length_mean = alpha * self.diag_length_mean + \
                (1 - alpha) * diag_length
        self.diag_length_list.append(diag_length)

    def compute_moving_distance(self, H):
        if len(self.position_list) <= 1:
            return 0
        # 如果无人机静止不动，
        # 车辆的开始位置为前一帧跟踪框的中心
        # 车辆的结束位置为当前帧跟踪框的中心
        # 当无人机发生运动
        # 车辆的开始位置为前一帧跟踪框中心 乘以 光流的矩阵
        # 车辆的结束位置仍为当前帧跟踪框的中心
        start_x = (self.position_list[-2][0] + self.position_list[-2][2]) / 2
        start_y = (self.position_list[-2][1] + self.position_list[-2][3]) / 2
        end_x = (self.position_list[-1][0] + self.position_list[-1][2]) / 2
        end_y = (self.position_list[-1][1] + self.position_list[-1][3]) / 2

        start_x, start_y, _ = H @ np.array([start_x, start_y, 1])
        d = math.sqrt((end_x - start_x)**2 + (end_y - start_y)**2)
        return d

    def add_position_and_moving_distance(self, trk_box: np.ndarray,
                                         H: np.ndarray, frame_index: int):
        self.add_position(trk_box)
        d = self.compute_moving_distance(H)
        self.moving_distance_list.append(d)
        self.last_frame_index = frame_index

        Q = 16  # 累计满足Q帧再计算速度
        FPS = 30  # 帧率
        L_p = 4.8  # 车辆对角线实际长度
        if len(self.moving_distance_list) >= 16:
            d = sum(self.moving_distance_list[-16:])
            diag = self.diag_length_mean
            R = L_p / diag
            t = (Q - 1) * 1 / 30
            v = d * R / t  # 单位为 m/s
            v = v * 3600 / 1000  # 转换为 km/h
            self.speed_list.append(v)
